#include "cartographer/mapping/internal/2d/scan_matching/real_time_correlative_scan_matcher_2d.h"

#include <algorithm>
#include <cmath>
#include <functional>
#include <limits>

#include "Eigen/Geometry"
#include "cartographer/common/lua_parameter_dictionary.h"
#include "cartographer/common/math.h"
#include "cartographer/mapping/2d/probability_grid.h"
#include "cartographer/mapping/internal/2d/tsdf_2d.h"
#include "cartographer/sensor/point_cloud.h"
#include "cartographer/transform/transform.h"
#include "glog/logging.h"

namespace cartographer
{
    namespace mapping
    {
        namespace scan_matching
        {
            namespace
            {
                float ComputeCandidateScore(const TSDF2D &tsdf, const DiscreteScan2D &discrete_scan, int x_index_offset, int y_index_offset)
                {
                    float candidate_score = 0.f;
                    float summed_weight = 0.f;
                    for (const Eigen::Array2i &xy_index : discrete_scan)
                    {
                        const Eigen::Array2i proposed_xy_index(xy_index.x() + x_index_offset, xy_index.y() + y_index_offset);
                        const std::pair<float, float> tsd_and_weight = tsdf.GetTSDAndWeight(proposed_xy_index);
                        const float normalized_tsd_score =
                            (tsdf.GetMaxCorrespondenceCost() - std::abs(tsd_and_weight.first)) / tsdf.GetMaxCorrespondenceCost();
                        const float weight = tsd_and_weight.second;
                        candidate_score += normalized_tsd_score * weight;
                        summed_weight += weight;
                    }
                    if (summed_weight == 0.f)
                        return 0.f;
                    candidate_score /= summed_weight;
                    CHECK_GE(candidate_score, 0.f);
                    return candidate_score;
                }

                float ComputeCandidateScore(const ProbabilityGrid &probability_grid, const DiscreteScan2D &discrete_scan,
                                            int x_index_offset, int y_index_offset)
                {
                    float candidate_score = 0.f;
                    for (const Eigen::Array2i &xy_index : discrete_scan)
                    {
                        const Eigen::Array2i proposed_xy_index(xy_index.x() + x_index_offset, xy_index.y() + y_index_offset);
                        const float probability = probability_grid.GetProbability(proposed_xy_index);
                        candidate_score += probability;
                    }
                    candidate_score /= static_cast<float>(discrete_scan.size());
                    CHECK_GT(candidate_score, 0.f);
                    return candidate_score;
                }

            } // namespace

            RealTimeCorrelativeScanMatcher2D::RealTimeCorrelativeScanMatcher2D(const proto::RealTimeCorrelativeScanMatcherOptions &options)
                : options_(options)
            {
            }

            std::vector<Candidate2D> RealTimeCorrelativeScanMatcher2D::GenerateExhaustiveSearchCandidates(
                const SearchParameters &search_parameters) const
            {
                int num_candidates = 0;
                // 预估大小
                auto &bounds = search_parameters.linear_bounds;
                for (int scan_index = 0; scan_index != search_parameters.num_scans; ++scan_index)
                {
                    const int num_linear_x_candidates = (bounds[scan_index].max_x - bounds[scan_index].min_x + 1);
                    const int num_linear_y_candidates = (bounds[scan_index].max_y - bounds[scan_index].min_y + 1);
                    num_candidates += num_linear_x_candidates * num_linear_y_candidates;
                }
                std::vector<Candidate2D> candidates;
                candidates.reserve(num_candidates);

                for (int scan_index = 0; scan_index != search_parameters.num_scans; ++scan_index)
                {
                    for (int x_index_offset = bounds[scan_index].min_x; x_index_offset <= bounds[scan_index].max_x; ++x_index_offset)
                    {
                        for (int y_index_offset = bounds[scan_index].min_y;
                             y_index_offset <= search_parameters.linear_bounds[scan_index].max_y; ++y_index_offset)
                        {
                            candidates.emplace_back(scan_index, x_index_offset, y_index_offset, search_parameters);
                        }
                    }
                }
                CHECK_EQ(candidates.size(), num_candidates);
                return candidates;
            }

            // std::vector<Candidate2D> RealTimeCorrelativeScanMatcher2D::GenerateExhaustiveSearchCandidates(const SearchParameters
            // &search_parameters) const
            // {
            //     int num_candidates = 0;
            //     // 预估大小
            //     auto &bounds = search_parameters.linear_bounds;
            //     for (int scan_index = 0; scan_index != search_parameters.num_scans; ++scan_index)
            //     {
            //         const int num_linear_x_candidates = (bounds[scan_index].max_x - bounds[scan_index].min_x + 1);
            //         const int num_linear_y_candidates = (bounds[scan_index].max_y - bounds[scan_index].min_y + 1);
            //         num_candidates += num_linear_x_candidates * num_linear_y_candidates;
            //     }
            //     std::vector<Candidate2D>().swap(particle_filter_->candidates);
            //     particle_filter_->candidates.reserve(num_candidates);
            //     for (int scan_index = 0; scan_index != search_parameters.num_scans; ++scan_index)
            //     {
            //         for (int x_index_offset = bounds[scan_index].min_x; x_index_offset <= bounds[scan_index].max_x; ++x_index_offset)
            //         {
            //             for (int y_index_offset = bounds[scan_index].min_y;
            //                  y_index_offset <= search_parameters.linear_bounds[scan_index].max_y; ++y_index_offset)
            //             {
            //                 candidates.emplace_back(scan_index, x_index_offset, y_index_offset, search_parameters);
            //             }
            //         }
            //     }
            //     CHECK_EQ(candidates.size(), num_candidates);
            //     return candidates;
            // }

            double RealTimeCorrelativeScanMatcher2D::Match(const transform::Rigid2d &initial_pose_estimate,
                                                           const sensor::PointCloud &point_cloud, const Grid2D &grid,
                                                           transform::Rigid2d *pose_estimate)
            {
                CHECK(pose_estimate != nullptr);
                // transform
                const Eigen::Rotation2Dd initial_rotation = initial_pose_estimate.rotation();
                auto angleAxis = Eigen::AngleAxisf(initial_rotation.cast<float>().angle(), Eigen::Vector3f::UnitZ());
                auto transform = transform::Rigid3f::Rotation(angleAxis);
                const auto rotated_point_cloud = sensor::TransformPointCloud(point_cloud, transform);

                const SearchParameters search_parameters(options_.linear_search_window(), options_.angular_search_window(),
                                                         rotated_point_cloud, grid.limits().resolution());
                // 以初始值为基础，生成一些列的旋转后的点云
                const auto rotated_scans = GenerateRotatedScans(rotated_point_cloud, search_parameters);
                // 根据旋转点云坐标，生成点对应的栅格坐标
                Eigen::Translation2f translation(initial_pose_estimate.translation().x(), initial_pose_estimate.translation().y());
                const auto discrete_scans = DiscretizeScans(grid.limits(), rotated_scans, translation);
                // 在x和y的维度以一定的步长采样若干候选的位姿，这里的x和y轴偏移量是指栅格偏移数，不是具体的距离
                std::vector<Candidate2D> candidates = GenerateExhaustiveSearchCandidates(search_parameters);
                // 计算每个候选位姿的概率得分
                ScoreCandidates(grid, discrete_scans, search_parameters, &candidates);
                // 选出最佳得分的候选位姿作为返回结果
                const Candidate2D &best_candidate = *std::max_element(candidates.begin(), candidates.end());
                *pose_estimate = transform::Rigid2d({initial_pose_estimate.translation().x() + best_candidate.x,
                                                     initial_pose_estimate.translation().y() + best_candidate.y},
                                                    initial_rotation * Eigen::Rotation2Dd(best_candidate.orientation));
                return best_candidate.score;
            }

            void RealTimeCorrelativeScanMatcher2D::ScoreCandidates(const Grid2D &grid, const std::vector<DiscreteScan2D> &discrete_scans,
                                                                   const SearchParameters &search_parameters,
                                                                   std::vector<Candidate2D> *const candidates) const
            {
                for (Candidate2D &candidate : *candidates)
                {
                    switch (grid.GetGridType())
                    {
                    case GridType::PROBABILITY_GRID:
                        candidate.score =
                            ComputeCandidateScore(static_cast<const ProbabilityGrid &>(grid), discrete_scans[candidate.scan_index],
                                                  candidate.x_index_offset, candidate.y_index_offset);
                        break;
                    case GridType::TSDF:
                        candidate.score = ComputeCandidateScore(static_cast<const TSDF2D &>(grid), discrete_scans[candidate.scan_index],
                                                                candidate.x_index_offset, candidate.y_index_offset);
                        break;
                    }
                    candidate.score *=
                        std::exp(-common::Pow2(std::hypot(candidate.x, candidate.y) * options_.translation_delta_cost_weight() +
                                               std::abs(candidate.orientation) * options_.rotation_delta_cost_weight()));
                }
            }
        } // namespace scan_matching
    }     // namespace mapping
} // namespace cartographer
